import numpy as np
from Xlib import display, X
import cv2
import time
import os

class X11ScreenCapture:
    def __init__(self):
        self.disp = display.Display()
        self.screen = self.disp.screen()
        self.root = self.screen.root
        self.width = self.screen.width_in_pixels
        self.height = self.screen.height_in_pixels
        self.format = X.ZPixmap  # 使用32位格式
        self.bytes_per_pixel = 4
        self.theoretical_row_bytes = self.width * self.bytes_per_pixel
        
        # 16块分割参数（4x4网格）
        self.grid_rows = 4
        self.grid_cols = 4
        self.block_width = self.width // self.grid_cols
        self.block_height = self.height // self.grid_rows
        
        # 存储块的最新状态（内存中）
        self.latest_blocks = {}  # 格式: (i,j) -> 块图像
        self.block_dir = "screen_blocks"
        os.makedirs(self.block_dir, exist_ok=True)

    def get_mouse_info(self):
        """获取鼠标位置和指针图像（统一接口）"""
        # 获取位置
        pointer_data = self.root.query_pointer()
        x, y = pointer_data.root_x, pointer_data.root_y
        
        # 获取指针图像和热点
        try:
            cursor = pointer_data.cursor
            shape = cursor.get_shape()
            hot_x, hot_y = shape.hot_x, shape.hot_y
            w, h = shape.width, shape.height
            cursor_data = np.frombuffer(shape.data, dtype=np.uint8).reshape((h, w, 4))
        except:
            #  fallback指针
            cursor_data = np.zeros((20, 20, 4), dtype=np.uint8)
            cv2.circle(cursor_data, (10, 10), 5, (255, 255, 255, 255), -1)
            hot_x, hot_y = 10, 10
            
        return (x, y), (hot_x, hot_y), cursor_data

    def capture_frame(self):
        """捕获原始屏幕帧（不含鼠标，鼠标将在拼接时添加）"""
        try:
            raw = self.root.get_image(
                x=0, y=0,
                width=self.width, height=self.height,
                format=self.format,
                plane_mask=0xFFFFFFFF
            )
            
            # 处理行字节
            try:
                actual_row_bytes = raw.bytes_per_line
            except AttributeError:
                actual_row_bytes = len(raw.data) // self.height

            # 转换为BGR格式
            data = np.frombuffer(raw.data, dtype=np.uint8)
            frame = []
            for y in range(self.height):
                start = y * actual_row_bytes
                end = start + self.theoretical_row_bytes
                row = data[start:end]
                if len(row) < self.theoretical_row_bytes:
                    row = np.pad(row, (0, self.theoretical_row_bytes - len(row)), mode='constant')
                frame.append(row)
            
            frame = np.concatenate(frame).reshape((self.height, self.width, 4))
            return cv2.cvtColor(frame, cv2.COLOR_BGRA2BGR)
        except Exception as e:
            print(f"捕获帧失败: {e}")
            return None

    def split_into_blocks(self, frame):
        """分割帧为16块（不含鼠标，鼠标在拼接时单独处理）"""
        blocks = []
        for i in range(self.grid_rows):
            for j in range(self.grid_cols):
                y1 = i * self.block_height
                y2 = (i + 1) * self.block_height if i < self.grid_rows - 1 else self.height
                x1 = j * self.block_width
                x2 = (j + 1) * self.block_width if j < self.grid_cols - 1 else self.width
                
                block = frame[y1:y2, x1:x2]
                blocks.append((i, j, block))
        return blocks

    def update_blocks(self, current_frame):
        """更新变化的块（仅屏幕内容，不含鼠标）"""
        current_blocks = self.split_into_blocks(current_frame)
        changed = []
        
        for i, j, curr_block in current_blocks:
            if (i, j) not in self.latest_blocks:
                # 新增块
                self.latest_blocks[(i, j)] = curr_block.copy()
                changed.append((i, j))
                cv2.imwrite(os.path.join(self.block_dir, f"block_{i}_{j}.png"), curr_block)
            else:
                # 比较差异
                prev_block = self.latest_blocks[(i, j)]
                if np.mean(cv2.absdiff(curr_block, prev_block)) > 5.0:
                    self.latest_blocks[(i, j)] = curr_block.copy()
                    changed.append((i, j))
                    cv2.imwrite(os.path.join(self.block_dir, f"block_{i}_{j}.png"), curr_block)
        
        return changed

    def stitch_with_mouse(self):
        """拼接块并添加鼠标指针（关键修改：在拼接后动态添加鼠标）"""
        if not self.latest_blocks:
            return None
            
        # 1. 拼接块得到基础图像
        stitched = np.zeros((self.height, self.width, 3), dtype=np.uint8)
        for (i, j), block in self.latest_blocks.items():
            y1 = i * self.block_height
            y2 = (i + 1) * self.block_height if i < self.grid_rows - 1 else self.height
            x1 = j * self.block_width
            x2 = (j + 1) * self.block_width if j < self.grid_cols - 1 else self.width
            
            # 处理尺寸偏差
            block = cv2.resize(block, (x2 - x1, y2 - y1))
            stitched[y1:y2, x1:x2] = block
        
        # 2. 获取最新鼠标信息并绘制（每次拼接都重新绘制，确保鼠标不消失）
        (x, y), (hot_x, hot_y), cursor = self.get_mouse_info()
        cursor_h, cursor_w = cursor.shape[:2]
        
        # 计算绘制位置
        draw_x = x - hot_x
        draw_y = y - hot_y
        
        # 计算有效绘制区域
        start_y = max(0, draw_y)
        end_y = min(stitched.shape[0], draw_y + cursor_h)
        start_x = max(0, draw_x)
        end_x = min(stitched.shape[1], draw_x + cursor_w)
        
        # 鼠标图像对应区域
        cursor_y1 = max(0, -draw_y)
        cursor_y2 = cursor_y1 + (end_y - start_y)
        cursor_x1 = max(0, -draw_x)
        cursor_x2 = cursor_x1 + (end_x - start_x)
        
        # 混合透明度
        alpha = cursor[cursor_y1:cursor_y2, cursor_x1:cursor_x2, 3] / 255.0
        rgb = cursor[cursor_y1:cursor_y2, cursor_x1:cursor_x2, :3]
        
        # 叠加鼠标到拼接图像上
        stitched[start_y:end_y, start_x:end_x] = (
            (1 - alpha[..., None]) * stitched[start_y:end_y, start_x:end_x] +
            alpha[..., None] * rgb
        ).astype(np.uint8)
        
        return stitched

    def close(self):
        self.disp.close()

def screen_video_stream():
    capture = X11ScreenCapture()
    print(f"开始捕获视频流（分辨率：{capture.width}x{capture.height}）")
    print(f"4x4块分割，鼠标将实时叠加在拼接图像上")
    
    try:
        while True:
            # 1. 捕获屏幕内容（不含鼠标）
            frame = capture.capture_frame()
            if frame is None:
                time.sleep(0.1)
                continue
            
            # 2. 更新变化的块（仅屏幕内容）
            changed = capture.update_blocks(frame)
            if changed:
                print(f"更新的块: {changed}")
            
            # 3. 拼接块并添加最新鼠标位置（每次都重新绘制鼠标）
            final_image = capture.stitch_with_mouse()
            if final_image is not None:
                cv2.imshow("Stitched Screen + Mouse (Press 'q' to exit)", final_image)
            
            # 退出检测
            if cv2.waitKey(1) & 0xFF == ord('q'):
                break
            time.sleep(0.05)
            
    finally:
        capture.close()
        cv2.destroyAllWindows()
        print("视频流已停止")

if __name__ == "__main__":
    screen_video_stream()